import random
import math
import plannersupport
import kernel
import grid
from bisect import bisect_left

class Diffusion:
    def __init__(self,settings,init_state):
        """
        Initializes the diffusion technique with the given settings and
        a state description.

        Four types of diffusion technique are currently supported:
        random, grid, kernel, and voronoi.

        Settings is a PlannerSettings structure, init_state is a state
        dictionary object"""
        
        self.settings = settings
        self.init_state = init_state
        #setup internal data structures
        if settings.diffusionTechnique=='grid' or settings.diffusionTechnique=='kernel':
            self.densityEstimators = []
            #set up the space discretization
            n = settings.diffusionDimensions
            numslices = int(math.ceil(float(len(self.vector_repr(init_state)))/n))
            self.diffusion_indices = [(i*n,i*n+n) for i in xrange(numslices)]
            for index in xrange(len(self.diffusion_indices)):
                width = 1.0
                if isinstance(settings.scaling,dict):
                    width = [s*settings.bandwidth for s in self.diffusion_repr(index,settings.scaling)]
                elif isinstance(settings.scaling,(int,float)):
                    width = float(settings.scaling)*settings.bandwidth
                else:
                    raise RuntimeError("Invalid settings.scaling variable")
                
                print 'Initializing diffusion',index,'width',width,'range',self.diffusion_indices[index]
                if settings.diffusionTechnique=='grid':
                    self.densityEstimators.append(grid.GridDensityEstimator(width))
                elif settings.diffusionTechnique=='kernel':
                    self.densityEstimators.append(kernel.AutoDensityEstimator(width))
        elif settings.diffusionTechnique=='voronoi':
            self.bmin = self.vector_repr(init_state)
            self.bmax = self.vector_repr(init_state)
            for index in xrange(len(self.bmin)):
                width = 1.0
                if isinstance(settings.scaling,dict):
                    width = vector_repr(settings.scaling)[index]*settings.bandwidth
                elif isinstance(settings.scaling,(int,float)):
                    width = float(settings.scaling)*settings.bandwidth
                else:
                    raise RuntimeError("Invalid settings.scaling variable")
                self.bmin[index] -= width
                self.bmax[index] += width
            self.nodes = []
        else:
            self.nodes = []

    def add_node(self,state,edge):
        if hasattr(self,'densityEstimators'):
            for i in xrange(len(self.densityEstimators)):
                self.densityEstimators[i].add(self.diffusion_repr(i,state),edge)
        elif hasattr(self,'bmin'):  #dynamic sampling region -- double on expansion
            self.nodes.append((state,edge))
            vstate = self.vector_repr(state)
            for i in xrange(len(self.bmin)):
                if vstate[i] < self.bmin[i]:
                    self.bmin[i] -= (self.bmax[i]-self.bmin[i])*0.5
                elif vstate[i] > self.bmax[i]:
                    self.bmax[i] += (self.bmax[i]-self.bmin[i])*0.5
        else:
            self.nodes.append((state,edge))

    def remove_node(self,state,edge):
        if hasattr(self,'densityEstimators'):
            for i in xrange(len(self.densityEstimators)):
                self.densityEstimators[i].remove(self.diffusion_repr(i,state),edge)
        else:
            delindex = (i for (i,p) in enumerate(self.nodes) if p[1]==edge).next()
            del self.nodes[delindex]

    def vector_repr(self,state):
        """Returns a vector-valued representation of the state."""
        newC = []
        for n in self.init_state:
            if n not in self.settings.irrelevant:
                newC.extend(plannersupport.flatten(state[n]))
        return newC

    def diffusion_repr(self,index,state):
        """Returns a vector-valued representation for the diffusion structure
        with index 'index'."""
        a,b = self.diffusion_indices[index]
        return self.vector_repr(state)[a:b]

    def pick_nodes(self,k,weight=None,cond=None):
        """Picks k nodes using the diffusion technique, with an optional
        weight/condition"""
        if self.settings.diffusionTechnique == 'grid':
            return self.pick_nodes_grid(k,weight,cond)
        elif self.settings.diffusionTechnique == 'kernel':
            return self.pick_nodes_kernel(k,weight,cond)
        elif self.settings.diffusionTechnique == 'voronoi':
            #NOTE: not implemented yet.
            return self.pick_nodes_voronoi(k,weight,cond)
        else:
            return self.pick_nodes_random(k,weight,cond)

    def pick_nodes_random(self,k,weight,cond):
        res = []
        if weight:
            w = []
            wsum = 0.0
            for (point,e) in self.nodes:
                wsum += weight(e)
                w.append(wsum)
            return self.nodes[bisect_left(w,random.uniform(0.0,wsum))][1].get_tail()
        elif cond:
            c = [e for (point,e) in self.nodes if cond(e)]
            assert(len(c) > 0)
            for i in xrange(k):
                res.append(random.choice(c).get_tail())
            return res
        else:
            return self.pick_nodes_random_raw(k)

    def pick_nodes_grid(self,k,weight,cond):
        """Picks k nodes at random using the grid structure. An O(m+kn/m)
        operation where m is the number of occupied grid cells.
        """
        if weight:
            v = sum([est.points.values() for est in self.densityEstimators],[])
            w = []
            wsum = 0.0
            for cell in v:
                wsum += sum(weight(e) for (point,e) in cell)/len(cell)
                w.append(wsum)
            res = []
            for i in xrange(k):
                cell = v[bisect_left(w,random.uniform(0.0,wsum))]
                #TODO: this could be faster if complete edges were stored
                #separately from active edges
                wc = []
                wcsum = 0.0
                for (point,e) in cell:
                    wcsum += weight(e)
                    wc.append(wcsum)
                res.append(cell[bisect_left(wc,random.uniform(0.0,wcsum))][1].get_tail())
            return res
        elif cond:
            v = sum([est.points.values() for est in self.densityEstimators],[])
            w = []
            wsum = 0.0
            for cell in v:
                if any (cond(e) for (point,e) in cell):
                    wsum += 1.0/float(len(cell))
                    #wsum += 1.0
                w.append(wsum)
            res = []
            for i in xrange(k):
                cell = v[bisect_left(w,random.uniform(0.0,wsum))]
                #TODO: this could be faster if complete edges were stored
                #separately from active edges
                res.append(random.choice([e for (point,e) in cell if cond(e)]).get_tail())
            return res
        else:
            return self.pick_nodes_grid_raw(k)
            
    def pick_nodes_kernel(self,k,weight,cond):
        """Picks k nodes at random using a kernel density estmator.  This is
        an O(n+klogn) operation.  However, the kernel estimator takes O(n)
        to add a point to the structure, so this is usually quite slow
        """
        #this is for a kernel density estimator
        if weight:
            est = random.choice(self.densityEstimators)
            assert min(est.density) > 0.0
            dsum = sum(est.density)
            w = []
            wsum = 0.0
            for (d,e) in zip(est.density,est.ids):
                #dsum/d is the weight for entry i
                wsum += weight(e)*dsum/d
                w.append(wsum)
            assert wsum > 0.0
            #weighted choice
            res = []
            for i in xrange(k):
                index = bisect_left(w,random.uniform(0.0,wsum))
                res.append(est.ids[index].get_tail())
            return res
        elif cond:
            est = random.choice(self.densityEstimators)
            assert min(est.density) > 0.0
            dsum = sum(est.density)
            ce = []
            w = []
            wsum = 0.0
            for (d,e) in zip(est.density,est.ids):
                if cond(e):
                    #dsum/d is the weight for entry i
                    wsum += dsum/d
                    w.append(wsum)
                    ce.append(e)
            assert wsum > 0.0
            #weighted choice
            res = []
            for i in xrange(k):
                index = bisect_left(w,random.uniform(0.0,wsum))
                res.append(ce[index].get_tail())
            return res
        else:
            return self.pick_nodes_kernel_raw(k)

    def pick_nodes_random_raw(self,k):
        res = []
        assert(len(self.nodes) > 0)
        for i in xrange(k):
            res.append(random.choice(self.nodes)[0])

    def pick_nodes_grid_raw(self,k):
        """Picks a node at random using the grid structure. An O(m+kn/m)
        operation where m is the number of occupied grid cells.
        """
        #this is for a grid density estimator
        #est = random.choice(self.densityEstimators)
        #assert(len(est.points) > 0)
        #v = est.points.values()
        v = sum([est.points.values() for est in self.densityEstimators],[])
        w = []
        wsum = 0.0
        for cell in v:
            if len(cell)>0:
                wsum += 1.0/float(len(cell))
            w.append(wsum)
        res = []
        for i in xrange(k):
            cell = v[bisect_left(w,random.uniform(0.0,wsum))]
            res.append(random.choice(cell)[1].get_tail())
        return res

    def pick_nodes_kernel_raw(self,k):
        """Picks k nodes at random using a kernel density estmator.  This is
        an O(n+klogn) operation.  However, the kernel estimator takes O(n)
        to add a point to the structure, so this is usually quite slow
        """
        #this is for a kernel density estimator
        est = random.choice(self.densityEstimators)
        assert min(est.density) > 0.0
        dsum = sum(est.density)
        w = []
        wsum = 0.0
        for (d,e) in zip(est.density,est.ids):
            #dsum/d is the weight for entry i
            wsum += dsum/d
            w.append(wsum)
        assert wsum > 0.0
        #weighted choice
        res = []
        for i in xrange(k):
            index = bisect_left(w,random.uniform(0.0,wsum))
            res.append(est.ids[index].get_tail())
        return res
